"""
Author        yongfa
Date          2023-09-04 05:37:52
LastEditTime  2023-09-07 03:23:01
LastEditors   yongfa
Description   邮箱登录类
"""
import os
import time
from email.parser import Parser
from email.header import decode_header
from email.utils import parseaddr
from email import message_from_string
import poplib
import imaplib
import re

from flask import current_app

from app.app_celery import logger
from app.extensions.utils.enums import ResumeEmailProtocolEnum
from app.extensions.exception import EMailLoginError, EMailProtocolError


class EmailLogin():

    def __init__(self, host: str, port: int, user: str, password: str, protocol: str, ssl: bool = False):

        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.protocol = protocol
        self.ssl = ssl
        self.mail_file_prefix = 'emails'
        self.upload_folder = current_app.config.get('UPLOAD_FOLDER')
        self.alllow_file_exts = ["doc", "docx", "pdf", "html", "xsl", "png", "jpg", "xlsx"]

    def login_pop3(self):
        try:
            if self.ssl:
                self.port = 995 if self.port is None else self.port
                conn = poplib.POP3_SSL(self.host, self.port)
            else:
                self.port = 110 if self.port is None else self.port
                conn = poplib.POP3(self.host)
            conn.set_debuglevel(1)
            conn.user(self.user)
            conn.pass_(self.password)
            logger.info("login host(%s) success", self.host)
            logger.info('Messages: %s. Size: %s' % conn.stat())
        except Exception as e:
            msg = "failed to login host(%s), username(%s)" % (self.host, self.user)
            logger.exception("%s, ; error:%s", msg, e)
            conn = None

        return conn

    def login_imap(self):
        try:
            if self.ssl:
                self.port = 993 if self.port is None else self.port
                conn = imaplib.IMAP4_SSL(self.host)
            else:
                self.port = 143 if self.port is None else self.port
                conn = imaplib.IMAP4(self.host)
            conn.login(self.user, self.password)
            logger.info("login host(%s) success", self.host)
            # logger.info('Messages: %s. Size: %s' % conn.status())
        except Exception as e:
            msg = "failed to login host(%s), username(%s)" % (self.host, self.user)
            logger.exception("%s, ; error:%s", msg, e)
            conn = None
        return conn

    def check_login(self):
        login_class = {
            ResumeEmailProtocolEnum.POP3.value: self.login_pop3,
            ResumeEmailProtocolEnum.IMAP.value: self.login_imap,
        }
        func = login_class.get(self.protocol)
        status = False
        if func is None:
            logger.error('protocol 不支持')
            return status
        conn = func()
        if conn is not None:
            status = True
        return status

    def get_mails(self, top_num: int = 1):
        email_class = {
            ResumeEmailProtocolEnum.POP3.value: self._get_mails_by_pop3,
            ResumeEmailProtocolEnum.IMAP.value: self._get_mails_by_imap,
        }

        func = email_class.get(self.protocol)
        if func is None:
            logger.error('protocol 不支持')
            return EMailProtocolError()
        result = func(top_num=top_num)
        return result

    def _get_mails_by_pop3(self, top_num: int = 1):
        """POP3 读取邮件"""
        conn = self.login_pop3()
        if conn is None:
            return EMailLoginError()
        # list()返回所有邮件的编号:
        resp, mails, octets = conn.list()

        index = len(mails)
        logger.info('get_mails:%s', index)
        result = []
        for i in range(index, index - top_num, -1):
            resp, lines, octets = conn.retr(i)
            msg_content = b'\r\n'.join(lines).decode('utf-8')

            # 解析邮件
            msg = Parser().parsestr(msg_content)
            data = self._parser_mail(msg, octets)
            result.append(data)

        conn.quit()
        return result

    def _get_mails_by_imap(self, top_num: int = 1):
        """IMAP 读取邮件"""
        conn = self.login_imap()
        if conn is None:
            return EMailLoginError()
        # list()返回所有邮件的编号:
        conn.select()
        typ, data = conn.search(None, 'ALL')
        mails = data[0].split()
        index = len(mails)
        logger.info('get_mails:%s', index)
        result = []

        for i in range(index, index - top_num, -1):
            typ, data = conn.fetch(mails[i - 1], '(RFC822)')
            # 使用utf-8 解码
            # 邮件编号
            mail_id = i - 1
            info = data[0][0].decode('utf-8').split(' ')[-1]
            _info = re.findall(r'\d+\.?\d*', info)
            if info:
                mail_id = int(_info[0])
            # 邮件内容
            content = data[0][1].decode('utf8')
            # 转化为email.message对象
            msg = message_from_string(content)

            data = self._parser_mail(msg, mail_id)
            result.append(data)

        conn.close()
        conn.logout()
        return result

    def _parser_mail(self, msg: str, mail_id: str):
        """解析邮件"""
        data = {'mail_id': mail_id}
        indent = 0
        for header in ['From', 'To', 'Subject', "Date"]:
            value = msg.get(header, '')
            if value:
                if header == 'Subject':
                    value = self._decode_str(value)
                elif header == 'Date':
                    _date1 = time.strptime(msg.get("Date")[0:24], '%a, %d %b %Y %H:%M:%S')
                    value = time.strftime("%Y-%m-%d %H:%M:%S", _date1)
                else:
                    hdr, addr = parseaddr(value)
                    name = self._decode_str(hdr)
                    value = u'%s <%s>' % (name, addr)
            data[header] = value
            logger.info('%s%s: %s' % ('  ' * indent, header, value))

        mail_file_path = "{}/{}".format(self.user, mail_id)
        content, attachments = self._save_att_file(msg, mail_file_path)
        data['content'] = content
        data['attachments'] = attachments
        return data

    def _decode_str(self, str_in: str):
        """字符编码转换"""
        value, charset = decode_header(str_in)[0]
        if charset:
            value = value.decode(charset)
        return value

    def _guess_charset(self, msg: str):
        """判断字符编码"""
        charset = msg.get_charset()
        if charset is None:
            content_type = msg.get('Content-Type', '').lower()
            pos = content_type.find('charset=')
            if pos >= 0:
                charset = content_type[pos + 8:].strip()
        return charset

    def _save_att_file(self, msg: str, mail_file_path: str):
        """附件下载函数"""

        content = ''
        attachment_files = []
        indent = 0
        mail_upload_path = os.path.join(self.mail_file_prefix, mail_file_path)
        save_path = os.path.join(self.upload_folder, mail_upload_path)
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type == 'text/plain' or content_type == 'text/html':
                content = part.get_payload(decode=True)
                charset = self._guess_charset(part)
                if charset:
                    content = content.decode(charset)
                # logger.info('%sText: \n%s' % ('  ' * indent, content + '...'))
            elif part.is_multipart():
                indent += 1
            else:
                file_name = part.get_filename()
                if file_name:
                    file_ext = file_name.split('.')[-1]
                    if file_ext in self.alllow_file_exts:
                        if not os.path.exists(save_path):
                            os.makedirs(save_path)
                        file_name = self._decode_str(file_name)
                        data = part.get_payload(decode=True)
                        file_path = os.path.join(save_path, file_name)
                        file_info = {
                            'name': file_name,
                            'path': file_path.replace(self.upload_folder, ''),
                        }
                        attachment_files.append(file_info)
                        with open(file_path, 'wb') as att_file:
                            att_file.write(data)
                        logger.info(f"附件 {file_name} 下载完成")
        return content, attachment_files
